(*
  Copyright (c) 2009 Barry Schwartz

  Permission is hereby granted, free of charge, to any person
  obtaining a copy of this software and associated documentation
  files (the "Software"), to deal in the Software without
  restriction, including without limitation the rights to use,
  copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the
  Software is furnished to do so, subject to the following
  conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  OTHER DEALINGS IN THE SOFTWARE.
*)

open XNum
open VMTypes
open UTypes
open SymbolTable
  
let true_sym = string_to_symbol (UString.uc_string_of_ascii "True")
let false_sym = string_to_symbol (UString.uc_string_of_ascii "False")
  
let make_tuple_pat elements =
  match elements with
  | [] -> assert false
  | [ x ] -> x
  | _ -> VMParser.PTuple elements
  
let compile_pattern off pat =
  let rec
    compile pat (((stack_depth, num_vars, var_names, checks) as state)) =
    match pat with
    | VMParser.PAnything ->
        (stack_depth, num_vars, var_names, (PCAnything :: checks))
    | VMParser.PId sym ->
        (stack_depth, (num_vars + 1), (sym :: var_names),
         ((PCVariable (off + num_vars)) :: checks))
    | VMParser.PNumber x ->
        (stack_depth, num_vars, var_names, ((PCNumber x) :: checks))
    | VMParser.PChar x ->
        (stack_depth, num_vars, var_names, ((PCChar x) :: checks))
    | VMParser.PSymbol sym ->
        (stack_depth, num_vars, var_names, ((PCSymbol sym) :: checks))
    | VMParser.PAssign (x, p) ->
        let (s, n, v, c) = compile p state
        in (s, (n + 1), (x :: v), ((PCAssign (off + n)) :: c))
    | VMParser.PTuple xs ->
        let len = List.length xs in
        let (s, n, v, c) = List.fold_right compile xs state
        in ((max ((s + len) - 1) stack_depth), n, v, ((PCTuple len) :: c))
    | VMParser.PList xs ->
        List.fold_right
          (fun x s1 ->
             let (s, n, v, c) = compile x s1
             in ((max (s + 1) stack_depth), n, v, (PCConsList :: c)))
          xs (stack_depth, num_vars, var_names, (PCNil :: checks))
    | VMParser.PListTail (xs, tail) ->
        List.fold_right
          (fun x s1 ->
             let (s, n, v, c) = compile x s1
             in ((max (s + 1) stack_depth), n, v, (PCConsList :: c)))
          xs (compile tail state) in
  let (stack_depth, num_vars, var_names, checks) = compile pat (0, 0, [], [])
  in (stack_depth, num_vars, (List.rev var_names), checks)
  
(* Creates a SymbolMap with all declared functions in |decls|. *)
let get_declared_variables decls =
  let rec iter names decls =
    match decls with
    | [] -> names
    | ((VMParser.DFun (name, pat, _, _) as d)) :: ds ->
        let defs = (try SymbolMap.find name names with | Not_found -> [])
        in
          ((match defs with
            | VMParser.DFun (_, p, _, _) :: _ -> (* check arity *)
                if (List.length p) <> (List.length pat)
                then
                  raise
                    (Syntax_error (Location.empty_location,
                       UString.append
                         (UString.uc_string_of_ascii
                            "arity mismatch in declaration of ")
                         (symbol_to_string name)))
                else ()
            | _ -> ());
           iter (SymbolMap.add name (d :: defs) names) ds)
    | VMParser.DPattern (p, _) :: ds ->
        let rec iter_pattern names p =
          (match p with
           | VMParser.PId name -> SymbolMap.add name [] names
           | VMParser.PAnything | VMParser.PNumber _ | VMParser.PChar _ |
               VMParser.PSymbol _ -> names
           | VMParser.PTuple ps -> List.fold_left iter_pattern names ps
           | VMParser.PList ps -> List.fold_left iter_pattern names ps
           | VMParser.PListTail (ps, p) ->
               List.fold_left iter_pattern (iter_pattern names p) ps
           | VMParser.PAssign (name, p) ->
               SymbolMap.add name [] (iter_pattern names p))
        in iter (iter_pattern names p) ds
  in iter SymbolMap.empty decls
  
let rec compile_expr scope expr =
  match expr with
    | VMParser.TUnbound -> [ BConst Unbound ]
    | VMParser.TId x ->
        (try [ Scope.lookup scope x ]
         with
           | Not_found ->
               raise
                 (Syntax_error (Location.empty_location,
                                UString.append
                                  (UString.uc_string_of_ascii "undefined symbol ")
                                  (symbol_to_string x))))
    | VMParser.TNumber x -> [ BConst (Number x) ]
    | VMParser.TChar x -> [ BConst (Char x) ]
    | VMParser.TFun cases -> compile_function scope cases
    | VMParser.TSymbol x ->
        if x = true_sym
        then [ BConst (Bool true) ]
        else
          if x = false_sym
          then [ BConst (Bool false) ]
          else [ BConst (Symbol x) ]
    | VMParser.TApp (f, args) ->
        (BApply (List.length args)) ::
          ((compile_expr scope f) @
             (List.fold_right (fun arg code -> (compile_expr scope arg) @ code)
                args []))
    | VMParser.TTuple xs ->
        let rec compile_elements is_const n els xs =
          (match xs with
             | [] -> (is_const, n, els)
             | y :: ys ->
                 let e = compile_expr scope y in
                 let c =
                   is_const && (match e with | [ BConst _ ] -> true | _ -> false)
                 in compile_elements c (n + 1) (e :: els) ys) in
        let (is_const, n, elements) = compile_elements true 0 [] xs
        in
          if is_const then
            let values = Array.make n (ref Unbound) in
            let _ =
              List.fold_left
                (fun i t ->
                   match t with
                     | [ BConst v ] -> (values.(i) <- ref v; i - 1)
                     | _ -> assert false)
                (n - 1) elements
            in
              [ BConst (Tuple values) ]
          else
            (BTuple n) :: (List.fold_left (fun es e -> e @ es) [] elements)
    | VMParser.TList xs ->
        let rec iter xs =
          (match xs with
             | [] -> [ BConst Nil ]
             | y :: ys ->
                 let t = compile_expr scope y in
                 let tail = iter ys
                 in
                   (match (t, tail) with
                      | ([ BConst a ], [ BConst b ]) ->
                          [ BConst (List (ref a, ref b)) ]
                      | _ -> BPair :: (t @ tail)))
        in iter xs
    | VMParser.TListTail (xs, tail) ->
        let rec iter xs =
          (match xs with
             | [] -> compile_expr scope tail
             | y :: ys ->
                 let t = compile_expr scope y in
                 let rest = iter ys
                 in
                   (match (t, rest) with
                      | ([ BConst a ], [ BConst b ]) ->
                          [ BConst (List (ref a, ref b)) ]
                      | _ -> BPair :: (t @ rest)))
        in iter xs
    | VMParser.TLocal (decls, expr) ->
        let (new_scope, init_code) = compile_local_declarations scope decls
        in BEndLocal :: ((compile_expr new_scope expr) @ init_code)
    | VMParser.TSequence (stmts, expr) ->
        (compile_expr scope expr) @
          (List.fold_left (fun code s -> (compile_statement scope s) @ code) []
             stmts)
    | VMParser.TDo exprs ->
        [ BFunction (1,
                     XList.rev_to_array
                       (BReturn ::
                          (List.fold_left (fun code s -> (compile_monad scope s) @ code)
                             [] exprs))) ]
    | VMParser.TIfThenElse (p, e0, e1) ->
        let then_code = compile_expr scope e0 in
        let else_code = compile_expr scope e1
        in
          else_code @
            (((BJump ((List.length else_code) + 1)) :: then_code) @
               ((BCondJump ((List.length then_code) + 2)) ::
                  (compile_expr scope p)))
    | VMParser.TMatch (expr, pats) ->
        (compile_matching scope
           (List.map (fun (p, g, e) -> ([ p ], g, e)) pats))
        @ (compile_expr scope expr)
and compile_matching scope cases =
  let num_pats = List.length cases in
  let (_, size, cases_code) =
    List.fold_left
      (fun (i, size, cases) (pats, g, e) ->
         let (stack_depth, num_vars, vars, compiled_pats) =
           List.fold_left
             (fun (depth, num_vars, vars, pats) pat ->
                let (sd, n, v, p) = compile_pattern num_vars pat
                in ((max depth sd), (num_vars + n), (v :: vars), (p :: pats)))
             (0, 0, [], []) pats in
           (* FIX: ensure that the variables are all distinct *)
         let (new_scope, _) =
           Scope.push scope (List.fold_left (fun a b -> b @ a) [] vars) in
         let expr_code = compile_expr new_scope e in
         let expr_size = List.length expr_code in
         let guard_code =
           match g with
             | None -> []
             | Some e ->
                 (BCondJump (expr_size + 3)) :: (compile_expr new_scope e) in
         let guard_size = List.length guard_code in
         let off =
           if i < num_pats then (guard_size + expr_size) + 4 else 0 in
         let cpat = XList.rev_to_array compiled_pats in
         let pop =
           if (Array.length cpat) = 1
           then BPop
           else BPopN (Array.length cpat) in
         let match_cmd =
           if (Array.length cpat) = 1
           then BMatch1 (cpat.(0), stack_depth, num_vars, off)
           else BMatchN (cpat, stack_depth, num_vars, off) in
         let code = expr_code @ ((pop :: guard_code) @ [ match_cmd ])
         in
           ((i + 1), (((size + expr_size) + guard_size) + 4),
            (code :: cases)))
      (1, 0, []) cases in
  let (_, code) =
    List.fold_right
      (fun case (size, code) ->
         let new_size = (size - (List.length case)) - 2
         in
           if new_size > 0
           then (new_size, (BEndLocal :: (BJump new_size) :: (case @ code)))
           else (new_size, (BEndLocal :: (case @ code))))
      cases_code (size, [])
  in code
and compile_function scope cases =
  let arity =
    match cases with | (ps, _, _) :: _ -> List.length ps | _ -> 0 in
  let rec make_dictionary dict cases =
    match cases with
      | [] -> Some dict
      | ([ VMParser.PSymbol s ], None, e) :: cs ->
          make_dictionary (SymbolMap.add s (compile_expr scope e) dict) cs
      | _ -> None
  in
    match make_dictionary SymbolMap.empty cases with
      | Some d ->
          let (code, syms) =
            SymbolMap.fold (fun s c (code, syms) -> ((c @ code), (s :: syms)))
              d ([], [])
          in (BDictionary (Array.of_list syms)) :: code
      | None ->
          [ BFunction (arity,
                       XList.rev_to_array (BReturn :: (compile_matching scope cases))) ]
and compile_local_declarations scope decls = (* extract a list of names *)
  let names = get_declared_variables decls in
  let vars = SymbolMap.fold (fun n _ vars -> n :: vars) names [] in
  let (new_scope, num_vars) = Scope.push scope vars in
    (* define the new symbols: first the functions, then the patterns *)
  let fun_defs =
    SymbolMap.fold
      (fun name defs code ->
         match defs with
           | [] -> code
           | (* a pattern, not a function *) _ ->
               let cases =
                 List.fold_left
                   (fun cases d ->
                      match d with
                        | VMParser.DFun (_, pats, guard, term) ->
                            (pats, guard, term) :: cases
                        | _ -> assert false)
                   [] defs in
               let (i1, i2) = Scope.lookup_local new_scope name
               in
                 (match cases with
                    | [ ([], None, term) ] ->
                        (BSet (i1, i2)) ::
                          ((compile_expr (Scope.shift scope 1) term) @ code)
                    | (* a variable declaration *) _ ->
                        (BSet (i1, i2)) ::
                          ((compile_function new_scope cases) @ code)))
      (* a real function        *) names [ BLocal num_vars ] in
  let rec iter code decls =
    match decls with
      | [] -> (new_scope, code)
      | VMParser.DFun (_, _, _, _) :: ds -> iter code ds
      | VMParser.DPattern (pat, term) :: ds ->
          (match pat with
             | VMParser.PId id -> (* Optimise the most common case. *)
                 let (i1, i2) = Scope.lookup_local new_scope id
                 in
                   iter
                     ((BSet (i1, i2)) :: ((compile_expr new_scope term) @ code))
                     ds
             | _ ->
                 let (sd, nv, v, c) = compile_pattern 0 pat in
                 let (local_scope, _) = Scope.push new_scope v in
                 let pat_code =
                   List.fold_left
                     (fun code var ->
                        let (a1, a2) = Scope.lookup_local new_scope var in
                        let (b1, b2) = Scope.lookup_local local_scope var
                        in (BSet (a1 + 1, a2)) :: (BVariable (b1, b2)) :: code)
                     [ BPop; BMatch1 (c, sd, nv, 0) ] v
                 in
                   iter
                     (BEndLocal ::
                        (pat_code @
                           ((compile_expr (Scope.shift scope 1) term) @ code)))
                     ds)
  in iter fun_defs decls
and compile_statement scope stmt =
  match stmt with
    | VMParser.SEquation (x, y) ->
        BUnify :: ((compile_expr scope x) @ (compile_expr scope y))
    | VMParser.SFunction t -> BPop :: (compile_expr scope t)
    | VMParser.SIfThen (p, s) ->
        let then_code = compile_statement scope s
        in
          then_code @
            ((BCondJump (List.length then_code)) :: (compile_expr scope p))
    | VMParser.SIfThenElse (p, s0, s1) ->
        let then_code = compile_statement scope s0 in
        let else_code = compile_statement scope s1
        in
          else_code @
            (((BJump (List.length else_code)) :: then_code) @
               ((BCondJump ((List.length then_code) + 1)) ::
                  (compile_expr scope p)))
and compile_monad scope stmt =
  match stmt with
    | VMParser.SEquation (x, y) ->
        BUnify :: ((compile_expr scope x) @ (compile_expr scope y))
    | VMParser.SFunction t -> (BApply 1) :: (compile_expr scope t)
    | VMParser.SIfThen (p, s) ->
        let then_code = compile_monad scope s
        in
          then_code @
            ((BCondJump ((List.length then_code) + 1)) ::
               (compile_expr scope p))
    | VMParser.SIfThenElse (p, s0, s1) ->
        let then_code = compile_monad scope s0 in
        let else_code = compile_monad scope s1
        in
          else_code @
            (((BJump ((List.length else_code) + 1)) :: then_code) @
               ((BCondJump ((List.length then_code) + 2)) ::
                  (compile_expr scope p)))
  
let rec compile_global_declarations scope decls =
  let names = get_declared_variables decls
  in
    ((* add the names to the global scope *)
     SymbolMap.iter (fun n _ -> Scope.add_global scope n Unbound) names;
     (* define the new symbols *)
     let fun_defs =
       SymbolMap.fold
         (fun name defs code ->
            match defs with
            | [] -> code
            | (* a pattern, not a function *) _ ->
                let cases =
                  List.fold_left
                    (fun cases d ->
                       match d with
                       | VMParser.DFun (_, pats, guard, term) ->
                           (pats, guard, term) :: cases
                       | _ -> assert false)
                    [] defs in
                let var = Scope.lookup_global scope name
                in
                  BUnify :: (BGlobal var) ::
                    ((compile_function scope cases) @ code))
         names [] in
     let rec iter code decls =
       match decls with
       | [] -> XList.rev_to_array code
       | VMParser.DFun (_, _, _, _) :: ds -> iter code ds
       | VMParser.DPattern (pat, term) :: ds ->
           let (sd, nv, v, c) = compile_pattern 0 pat in
           let (local_scope, _) = Scope.push scope v in
           let pat_code =
             List.fold_left
               (fun code var ->
                  let v = Scope.lookup_global scope var in
                  let (b1, b2) = Scope.lookup_local local_scope var
                  in BUnify :: (BGlobal v) :: (BVariable (b1, b2)) :: code)
               [ BMatch1 (c, sd, nv, 0) ] v
           in iter (pat_code @ ((compile_expr scope term) @ code)) ds
     in iter fun_defs decls)
  
let compile_declarations scope stream =
  let lexer = Lexer.make_lexer (Scope.symbol_table scope) stream in
  let decls = VMParser.parse_program lexer
  in compile_global_declarations scope decls
  
let compile_expression scope stream =
  let lexer = Lexer.make_lexer (Scope.symbol_table scope) stream in
  let expr = VMParser.parse_expression lexer
  in Array.of_list (List.rev (compile_expr scope expr))
  

